Constrained decoding with weighted automata1

I want to share a take on grammar-constrained generation I’ve been working on for a while. This post is a preliminary writeup: enough detail to communicate the core “read the (G)LR stack with a weighted automaton” idea (and timestamp it). I’ll likely follow up with a post on the various tricks you need to actually build the damn things without the compiler exploding.

The problem

You have a grammar (often a JSON Schema, or “JSON-but-with-some-fields”, or a programming-language subset) and you want the LLM to produce grammatically valid output.

At each step, the LLM outputs logits over the vocabulary $V$. You sample a token. Append it to the output. Repeat.

Constrained decoding is: don’t sample tokens that would make the output grammatically invalid.

So we maintain a constraint state and, for each step, compute a mask $M \subseteq V$ of allowed tokens, and apply it to the logits.

No more broken JSON.

A precise way to say “allowed” (this matters later):

A token $v \in V$ is valid given the current prefix $w$ if there exists some continuation $u$ such that $wvu$ is in the grammar language.

Equivalently, the mask is

That “$\exists u$” is the whole game: we’re not asking “does this token complete a valid document right now?”, we’re asking “could it still lead to a valid completion?”

Two primitive operations

Think of constrained decoding as two functions:

• get_mask(state) -> bitmask: returns a bitmask $M \subseteq V$ over the LLM vocabulary.
• commit(state, token) -> state: updates the constraint state with the chosen token.

At runtime the loop looks like:

loop:
  parallel:
    GPU: logits
    CPU: mask = get_mask(state)
  apply mask to logits
  sample token
  state = commit(state, token)

Nothing fancy, except: get_mask is on the critical path of decoding.

Why worst-case latency matters

LLM inference is batched: multiple user requests run on the GPU at once, and each step wants a mask at the same time.

If you want to generate ~1k tokens/sec, you have ~1ms to produce a mask. And if you miss even a single one—if a mask takes 2ms instead of 1—then in a batched setting it screws up scheduling: the GPU sits waiting on a straggler CPU mask, and suddenly your p99 latency is trash.

Masks can’t be late; they have to arrive on time every step. It’s worst-case performance that matters.

It’s not “can you do masking fast on average?”
It’s “can you do it fast every single step?”

Commit ~ incremental parsing

The commit op is basically incremental parsing: you’re incrementally parsing a growing prefix.

Many battle-tested systems like tree-sitter and Lezer use Generalized LR (GLR) parsing with a graph-structured stack (GSS).

• An LR parser maintains a stack of parser state IDs.
• It decides what to do by looking at the top of stack plus the next terminal symbol.
• The “Generalized” part handles ambiguity: instead of a single stack, you keep a graph of possible stacks sharing structure (the GSS). When ambiguity arises, stacks fork; when they converge to the same state, they merge; error paths are pruned.

So commit is not the “new idea” here. commit is: use a decent incremental GLR parser and keep it compact.

The interesting part is get_mask.

Glossary

LLM Token

An element of the model vocabulary—a BPE token with an ID and byte string. ~200k of these in modern tokenizers.

Grammar Terminal

The symbols your parser actually consumes. Might be bytes, characters, or lexer tokens like --, IDENT, STRING.

Parser State ID

An integer identifying the LR automaton's current state. The parser stack is a sequence of these.

Generating masks

Now you have a parse state (often a GSS), and you want a mask over LLM tokens.

A straightforward approach is:

For each LLM token $t \in V$, check whether appending it keeps the parse valid.

But you don’t append an LLM token to the parser. You append terminal symbols.

And a single LLM token is a byte string which might:

• contain multiple terminals ("}\n" could be RBRACE, NEWLINE, …),
• end partway through a terminal (halfway through a string literal, halfway through a number, halfway through true),
• or be lexically ambiguous depending on context.

So “try token $t$” is really “try every way $t$’s bytes could advance the lexer into some terminal sequence, and then try advancing the parser for that terminal sequence.”

That mismatch is where a lot of “obvious” masking schemes get wrecked.

Example: segmentation explosions

Here’s a toy example that captures the shape of the problem.

Suppose your terminals include - and -- (or more generally, terminals that overlap). If you treat tokenization as “split into terminals in any way that works” (scannerless style), then a run of $N$ dashes has $F_{N+1}$ possible segmentations (Fibonacci growth), because you’re tiling length $N$ using pieces of length 1 and 2.

And tokenizers absolutely contain long weird tokens. In cl200k there are tokens that are on the order of 100+ repeated punctuation characters. For $N \approx 112$, $F_{N+1}$ is $\sim 10^{23}$ segmentations.

That’s for a single candidate token.

And there are ~200k tokens to check.

This is obviously not something you can do per decoding step.

Why tries/trellises still hurt in p99.9 land

A common direction is:

• Build a trie/trellis representing how LLM tokens map to terminal sequences.
• Traverse that trie while advancing the parser state, pruning invalid paths.
• Collect the set of LLM tokens that survive.

This can work well for many grammars and workloads. People implement this and get decent average performance, especially for “simple” grammars.

But for general CFG-ish constraints (especially when you care about worst-case), you run into two issues:

1) You end up doing real parser work “inside” masking

Even if your lexical side is deterministic (say you enforce longest-match, or you have a DFA lexer), a single grammar terminal shift can trigger a chain of reductions.

In an LR parser, “shift terminal $t$” is often really:

1. while top-of-stack wants to reduce, do some reductions (pop some states, push a goto state),
2. then shift $t$.

Those reduction chains are table-driven and fast in the happy path, but in GLR they can fan out. The operations are pointer-heavy and cache-unfriendly: you’re manipulating a GSS, merging nodes, pruning branches, etc.

If you do that inside masking—i.e. while traversing a trie of possible tokens—you’re effectively building and mutating a GSS for a hypothetical next token, and doing it potentially thousands of times per step.

That’s exactly the kind of p99.9 work you don’t want on the critical path.

2) You still pay for long/ugly tokens

Even with tries, there are tokens whose byte strings correspond to long sequences of terminals (think: long runs of punctuation, or tokens containing lots of structural characters plus whitespace/newlines). Those paths can survive pruning for a while, so you end up doing a lot of terminal-level processing per mask computation.

And yes, long ugly sequences can be syntactically valid in real languages. Python, JS, etc. happily accept monstrosities like:

------------------1

That means “a lot of unary minus operators.” It’s valid, and it forces the parser to do real work on a long terminal sequence.

You can transform the grammar to remove unit/null reductions. You can aggressively merge equivalent GSS nodes. You can add clever memoization.

I tried hard to make the “trie + incremental parse simulation” approach behave well in worst-case latency terms. In my experience, it’s a dead end if you’re aiming for predictable sub-millisecond masking on arbitrary inputs/grammars.

Invert the problem

Token validity depends on what’s on the stack.

Instead of asking:

“does LLM token $t$ work on this stack?” (for each $t$)

ask:

“given this stack, which LLM tokens work?”

That’s the reframing.

Now we want a data structure that turns:

stack → allowed-token bitset

into something we can execute quickly.

That’s where the weighted automaton comes in.

Weighted automata over parser states

Think of a finite automaton, except:

• each transition carries a weight (here: a bitset over tokens), and
• when multiple paths lead to the same state, you combine their weights.

In the automata I’m talking about:

• Input symbols are LR parser state IDs (i.e. elements of the parse stack).
• Weights are bitsets over the LLM vocabulary $V$.
• Traversing a transition applies an intersection (filter tokens).
• Merging paths onto one state applies a union (tokens valid via any path).

So it’s like “run a DFA on the stack”, but with “which tokens survive” flowing through it.

The automaton reads a representation of your current parse configuration:

• For a single LR stack: the sequence of state IDs from top to bottom.
• For GLR: a GSS, i.e. a DAG representing many possible stacks sharing structure.

A token is valid if it’s valid on any parse path in the GSS, so on a GSS you just union the results across paths.

Another way to see it (why weights help)

Another mental model that helped me:

1. Fix an LLM token $t$.
2. Consider all the ways its bytes could be tokenized into grammar terminals (given incremental lexing, overlapping terminals, etc.).
3. For each such terminal sequence, consider all stack configurations from which that sequence can be legally parsed (shift/reduce/goto) without error.

That set of stacks is a regular language over parser state IDs (this is closely related to the classical “viable prefixes are regular” result from LR theory).

So you can imagine an automaton $A_t$ that recognizes “stacks on which token $t$ is valid.”

Of course, if you do that for every $t \in V$, you’d have 200k automata. That’s useless at runtime.

A weighted automaton is how you smash those 200k membership tests into one run:

• Transitions are annotated with “which tokens would accept if we take this transition.”
• Running it once on the stack gives you exactly: “which tokens’ $A_t$ would accept this stack.”

Same computation, but vectorized across tokens via bitsets.

A nice optimization: stop reading the stack early

If you implement the obvious “read stack symbols until you hit bottom” approach, you’ll still sometimes do more work than needed: often, after reading a small suffix of the stack, the set of valid tokens is already determined and deeper stack symbols can’t change it.

You can bake this into the automaton:

• Each automaton state has an “already-valid” contribution (tokens that are valid regardless of what comes next).
• Outgoing transitions only carry the “still-unknown” contribution.

Operationally, it’s the difference between:

• “keep filtering ALL tokens forever”
• vs
• “as soon as some tokens are decided, peel them off, and only keep propagating the remainder”

So you can stop when there’s nothing left in flight.

I’m not going to cite code here, but conceptually it’s just a precomputed fixed-point thing: for each automaton state, compute which tokens are guaranteed from this point (“final weight”), subtract them from the weights you propagate forward, and you get a monotone decreasing set that hits empty quickly.

Runtime sketch (single stack)

For a single LR stack, the runtime looks like:

def get_mask(stack_state_ids_top_to_bottom):
    # frontier: map automaton_state -> bitset(tokens)
    frontier = { A.start: ALL_TOKENS }

    decided = EMPTY  # tokens already known valid from suffix read so far

    for sid in stack_state_ids_top_to_bottom:
        new = {}
        for a_state, tokens in frontier.items():
            for (a2, weight) in A.step(a_state, sid):
                tokens2 = tokens & weight                 # intersection (filter)
                if tokens2.any():
                    new[a2] = new.get(a2, EMPTY) | tokens2  # union (merge paths)
        frontier = new

        # Optionally accumulate “final/decided” contribution here and
        # remove it from frontier tokens, so frontier only contains
        # “still-undecided” tokens.

        if not frontier:               # nothing left that depends on deeper stack
            break

    return decided | combine_accepting(frontier)

This is the key operational point:

• You are not iterating over the vocabulary.
• You are not simulating big chains of reductions at runtime.
• You are doing:
  • transition-table lookups (CPU-cache-friendly),
  • bitset intersections,
  • bitset unions at merges.

Runtime on a GSS

For a GSS, you do basically the same computation, but over a graph rather than a single list.

Conceptually:

• each edge in the GSS is labeled by a parser state ID,
• each node represents a shared suffix,
• you propagate $\cap$/$\cup$ through the product of:
  • GSS structure (many stacks),
  • automaton transitions.

A token is valid if it’s valid on any stack path, so whenever GSS paths merge you union their token sets, and whenever automaton paths merge you union there too.

The important part is: it’s still the same “bitset flows through a graph via ∩ and ∪” pattern. No backtracking, no “try token $t$” loop.

So where does the automaton come from?

Up to now I’ve treated “the weighted automaton” as magic.

It’s not magic, it’s just precomputation.

At a high level, you want an automaton that answers:

given the current lexer+parser state (represented by the stack/GSS), which LLM tokens could lead to a valid continuation?

There are two distinct problems mixed together:

1. Lexical: what grammar terminals could a given LLM token’s bytes produce (given where we are in the lexer)?
2. Syntactic: if we fed those terminals to the LR parser, would they be legal given the current stack?

The thing I compile is basically a composition of two automata:

• a token→terminal-sequences automaton (“Terminal DWA”),
• and a terminal→stack-effect automaton (“Template automata”), composed into a final Parser DWA that reads stack state IDs and outputs valid-token bitsets directly.

I’ll keep this section high-level (the details are where the compile-time pain is), but the decomposition matters because it explains why the runtime is so simple.

1) Token → terminals (Terminal DWA)

An LLM token is a byte string.

A grammar tokenizer/lexer consumes a stream of bytes and emits grammar terminals. Crucially:

• the lexer may be in the middle of recognizing a terminal,
• and longest-match means “a match might be extendable,” so incremental lexing can require keeping multiple active lexer configurations (more on this later).

So the mapping “token $v$ → terminal sequence” is not a single fixed lookup. It depends on the current lexer state.

The Terminal DWA is the precomputed structure that answers:

from lexer state $x$, which LLM tokens $v$ can produce which terminal $t$ next (and what lexer state do we end up in)?

A practical way to build it is:

• build a trie over all vocabulary token byte strings,
• simulate the grammar lexer over that trie,
• record (state, terminal, next_state) edges weighted by “which vocabulary tokens realize this edge.”

Then determinize.

The key thing the Terminal DWA buys you is: it collapses “iterate over 200k tokens and run the lexer” into “traverse a small automaton state space and get bitsets.”

2) Terminal → stack effect (Template automata)

Now suppose the lexer says “the next terminal is $t$.” What does the LR parser do?

• Sometimes it just shifts $t$.
• Often it does a chain of reductions first (pop a bunch of states, push goto states), and then shifts $t$.
• Those reduction chains depend on the top of the stack, but importantly: they depend on a bounded suffix of the stack in well-behaved grammars.

You can precompute, for each terminal $t$, an automaton that reads stack state IDs (top down) and represents all the ways the parser could legally process $t$, including the reduction chain.

Internally I represent the net stack transformation using a push/pop algebra (the polycyclic monoid is the clean formalism), but you don’t need to care about that to get the intuition:

• the template automaton for terminal $t$ encodes “how many states might get popped before we can shift $t$, and what goto transitions get pushed as a result.”

This is the piece that takes “pointer-heavy GLR reduction simulation” off the runtime path and turns it into precomputed transitions.

3) Compose them into a Parser DWA

Finally, you compose:

• Terminal DWA edges (which say “token $v$ could yield terminal $t$ from lexer state $x$”) with
• Template automata (which say “terminal $t$ is legal from stack suffix $s$”)

to get one deterministic weighted automaton that:

• takes a lexer configuration + a parser stack suffix as input,
• and outputs the set of LLM tokens that can extend the parse.

This is the automaton you run in get_mask.

One subtle but very useful algebraic fact during composition: when you concatenate stack-effect templates, you get adjacent “push then pop” pairs that cancel. That lets you simplify aggressively while building the composed automaton (and prune impossible paths when pushes/pops mismatch).

Also: at runtime I don’t actually need the automaton to update the stack. commit is implemented by the real lexer+parser. get_mask just needs to know whether a token is valid, so the automaton can discard some “push” bookkeeping once it has done its filtering job.

Why you only read a bounded suffix of the stack

A question I always get is: “do you really scan the entire stack every time?”

No. In practice you only need a bounded suffix, and you can make that precise.

Two related reasons:

• In LR theory, the set of viable prefixes (i.e. possible stack configurations) is regular, which is why an automaton over stack states even makes sense here.
• If you normalize the grammar to eliminate certain patterns (right recursion / hidden left recursion in the reduction chains), then the reduction chains that matter for “can I shift terminal $t$?” are bounded. That makes the per-terminal template automata acyclic, which forces a bound on how far down the stack you ever need to look.

In implementation terms: masks stabilize quickly, and you can early-exit.

Tokenization ambiguity (and longest-match) in streaming generation

The term “token” is overloaded:

• The LLM tokenizer maps text to subword units from a fixed vocabulary (byte strings).
• The grammar tokenizer/lexer maps a byte stream to grammar terminals (IDENT, "{", NUMBER, keywords, etc.).

These tokenizations don’t align.

Even if your lexer uses longest-match, longest-match is inherently forward-looking: you can’t know a match is final until you see what comes next.

Classic example:

• "+" followed by "+" should yield INCREMENT ("++"), not two PLUS.

In a streaming setting, after you see the first "+", you’re in a state where:

• you have a valid match for PLUS,
• but it’s extendable to INCREMENT.

So you have to represent that uncertainty somehow.

The approach I use is what I call inhibited terminals (this is just a convenient operational trick, not a deep theory term):

• when a terminal $t$ matches but remains extendable, you fork:
  • one branch commits $t$ now (but marks it “inhibited”: it will be invalidated if a longer match appears),
  • another branch waits for more bytes.
• if later input yields a longer match, the premature branch is pruned.

This plays nicely with GLR+GSS, because “fork and prune” is already the parser’s native move.

Practically, the constraint state you carry around is not just “parser stack(s)”; it’s “(parser stack(s), lexer state(s))”. get_mask needs to condition on both, which is why the Parser DWA has multiple initial states (one per active lexer state).

Putting it together

At runtime you now have two different kinds of work, with very different performance profiles:

commit(state, token) -> state

• Uses GLR machinery.
• Feeds the token’s bytes into the grammar lexer incrementally.
• Whenever the lexer emits a terminal, apply LR shift/reduce on the GSS.
• Maintain compact ambiguity, prune errors, merge when states coincide.

This is the inherently “CFG-y” part. It can still have pathological cases (that’s life), but you at least get decades of parsing engineering behind it.

get_mask(state) -> mask

• Uses the precompiled weighted automaton (Parser DWA).
• Treats the current GSS/stack suffixes as the “input string(s)” over parser state IDs.
• Propagates token bitsets through the automaton using $\cap$ on transitions and $\cup$ on merges.
• Stops early when deeper stack symbols can’t change the result.
• Returns a vocabulary mask.

This is the performance-critical part, and it’s exactly where precomputation buys you predictability.

Why this feels close to optimal

I’m deliberately avoiding “this is optimal” as a theorem. Parsing isn’t a domain where you get many satisfying worst-case optimality results that also match real-world constraints.

But it feels close to optimal for two reasons:

On the commit side: GLR/GSS is a strong local optimum

For incremental CFG parsing with ambiguity:

• table-driven decisions,
• shared structure via GSS,
• pruning/merging,

…is a pretty hard baseline to beat. There’s a reason people working on these problems keep converging on GLR-like techniques:

• https://tratt.net/laurie/blog/2020/which_parsing_approach.html
• https://marijnhaverbeke.nl/blog/lezer.html
• https://tree-sitter.github.io/tree-sitter/creating-parsers/3-writing-the-grammar.html
• https://www2.eecs.berkeley.edu/Pubs/TechRpts/1997/CSD-97-946.pdf

You still inherit GLR’s lack of comforting worst-case bounds in “fully adversarial ambiguity” settings. In theory, GLR can degrade badly on highly ambiguous grammars/inputs. In practice, for “programming-language-ish” grammars with reasonable disambiguation, it behaves well.

On the mask side: “read once, never backtrack” is what you want

The weighted-automaton mask computation does the minimum work you’d reasonably hope for:

• scan a bounded portion of the stack once,
• do predictable inner-loop ops (bitset ∩ and ∪),
• stop as soon as the result is determined,
• never iterate over the vocabulary,
• never do speculative parser reductions for hypothetical tokens.

In other words: the runtime work scales with the size of the parse configuration, not with vocabulary size and not with “how gnarly are the tokens.”

Fast run, slow compile

The cost shifts to compile time.

Precompiling the grammar into this Parser DWA involves determinization and simplification in a semiring-ish setting (bitset weights, union at merges, intersection along paths). If you do it naively, large grammars can blow up in memory/time.

Getting compile-time and memory to behave took most of my engineering effort:

• determinization order matters a lot,
• when you push bitsets around you must be careful about sharing/copying,
• you need aggressive pruning/cancellation,
• you need to normalize certain grammar patterns to keep reduction chains bounded,
• and you really want a representation that plays well with CPU caches.

That’s probably the next post.

If you’re building constrained decoding and you care about p99.9 latency, my main takeaway is:

• don’t put real parsing work in get_mask,
• make commit do the parsing once for the chosen token,
• and make get_mask a precomputed “read stack → bitset” computation.

That inversion is the whole trick.